package ssau.knyazev.fft.modules.oob;

import ssau.knyazev.core.Complex;
import ssau.knyazev.fft.common.FFTConst;
import ssau.knyazev.modules.Generator;
import ssau.knyazev.modules.Helper;
import ssau.knyazev.modules.Printer;

public class FFTModule {
	
	protected FFTModule() {	}
	
	public static Complex[] dft1D(float[] src) {
		Complex[] arr = Helper.transformFloat(src);
		return dft1D(arr);
	}
	
	public static Complex[] dft1D(Complex[] src){
		return dft1D(src, false);
	}
	
	public static Complex[] inverseDFT1D(Complex[] src) {
		return dft1D(src, true);
	}
	
	public static Complex[] dft1D(Complex[] src, boolean inverse) {
		Complex[] res = new Complex[src.length];
		for (int k = 0; k < res.length; k++) {
			res[k] = new Complex();
			if (inverse){
				for (int i = 0; i < res.length; i++){
					res[k] = res[k].add(src[i].multy(Complex.W(res.length, -k * i)));
				}
				res[k] = res[k].devideBy(res.length);
			} else {
				for (int i = 0; i < res.length; i++){
					res[k] = res[k].add(src[i].multy(Complex.W(res.length, i * k)));
				}
			}
		}
		return res;
	}

	public static Complex[][] dft2D(float[][] src) {
		Complex[][] arr = Helper.transformFloat(src);
		return dft2D(arr);
	}
	
	public static Complex[][] dft2D(Complex[][] src){
		return dft2D(src, false);
	}

	public static Complex[][] inverseDFT2D(Complex[][] src) {
		return dft2D(src, true);
	}
	
	public static Complex[][] dft2D(Complex[][] src, boolean inverse) {
		Complex[][] res = new Complex[src.length][src[0].length];
		for (int i = 0; i < src.length; i++) {
			for (int j = 0; j < src[0].length; j++) {
				res[i][j] = new Complex();
				if (inverse){
					for (int m = 0; m < src.length; m++) {
						for (int n = 0; n < src[0].length; n++) {
							res[i][j] = res[i][j].add(src[m][n].multy(Complex.W(src.length, -m * i)).multy(Complex.W(src[0].length, -n * j)));
						}
					}
					res[i][j] = res[i][j].devideBy(res.length * res[0].length);
				} else {
					for (int m = 0; m < src.length; m++) {
						for (int n = 0; n < src[0].length; n++) {
							res[i][j] = res[i][j].add(src[m][n].multy(Complex.W(src.length, i * m)).multy(Complex.W(src[0].length, j * n)));
						}
					}
				}
			}
		}
		return res;
	}

	public static Complex[] fft1D(float[] src){
		Complex[] arr = Helper.transformFloat(src);
		return fft1D(arr);
	}
	
	public static Complex[] fft1D(Complex[] src){
		return fft1D(src, false);
	}
	
	public static Complex[] inverseFFT1D(Complex[] src){
		Complex[] res = fft1D(src, true);
		for (int i = 0; i < res.length; i++){
			res[i] = res[i].devideBy(res.length);
		}
		return res;
	}

	public static Complex[] fft1D(Complex[] src, boolean inverse){
		inverse(src);
		int N = src.length;
		int size = 0;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		Complex[] w_list = new Complex[src.length];
		if (inverse){
			for (int i = 0; i < w_list.length; i++){
				w_list[i] = Complex.W(N, -i);			
			}
		} else {
			for (int i = 0; i < w_list.length; i++){
				w_list[i] = Complex.W(N, i);			
			}
		}
		for (int q = 0; q < size; q++) {
			int step = (int) Math.pow(2, q);
			int step2 = step*2;
			boolean flag = true;
			int count = 0;
			for (int i = 0; i < N - step; i++) {
				if (flag) {
					Complex t = w_list[(i % (step2))*(N/step2)];
					butterfly(src[i], src[i + step], t);
				}
				count++;
				if (count == step) {
					flag = !flag;
					count = 0;
				}
			}
		}
		return src;
	}

	public static Complex[][] fft2D(float[][] src){
		return fft2D(src, FFTConst.SIMPLE_MODE);
	}
	
	public static Complex[][] fft2D(float[][] src, int mode){
		Complex[][] arr = Helper.transformFloat(src);
		return fft2D(arr, mode);
	}
	
	public static Complex[][] fft2D(Complex[][] src){
		return fft2D(src, FFTConst.SIMPLE_MODE);
	}
	
	public static Complex[][] inverseFFT2D(Complex[][] src){
		Complex[][] res = fft2D(src, FFTConst.SIMPLE_MODE, true);
		for (int i = 0; i < res.length; i++)
			for (int j = 0; j < res.length; j++){
				res[i][j] = res[i][j].devideBy(res.length*res.length);
			}
		return res;
	}
	
	public static Complex[][] fft2D(Complex[][] src, int mode){
		return fft2D(src, mode, false);
	}
	
	public static Complex[][] fft2D(Complex[][] src, int mode, boolean inverse){
		if (mode == FFTConst.ROWCOL_MODE){
			return fft2DRowCol(src);
		}
		inverse(src);
		int N = src.length;
		int size = 0;
		for (int x = src.length; x > 1; x = x >> 1){
			size++;
		}
		Complex[] w_list = new Complex[src.length];
		if (inverse){
			for (int i = 0; i < w_list.length; i++){
				w_list[i] = Complex.W(N, -i);			
			}
		} else {
			for (int i = 0; i < w_list.length; i++){
				w_list[i] = Complex.W(N, i);			
			}
		}
		for (int q = 0; q < size; q++) {
			int step = 1 << q;
			int step2 = step << 1;
			for (int down = 0; down < N - step; down++) {
				for (int right = 0; right < N - step; right++) {
					boolean flag = (down/step) % 2 + (right/step) % 2 > 0 ? false : true;
					if (flag) {
						Complex t1 = w_list[(down % step2)*(N/step2)];
						Complex t2 = w_list[(right % step2)*(N/step2)];
						butterfly(src[down][right], src[down][right+step], src[down+step][right], src[down+step][right+step], t1, t2);
					}
				}
			}		
		}
		return src;
	}
	
	protected static Complex[][] fft2DRowCol(Complex[][] src){
		for (int i = 0; i < src.length; i++){
			src[i] = fft1D(src[i]);
		}
		for (int i = 0; i < src.length; i++){
			Complex[] tarr = new Complex[src.length];
			for (int j = 0; j < src.length; j++){
				tarr[j] = src[j][i];
			}
			tarr = fft1D(tarr);
			for (int j = 0; j < src.length; j++){
				src[j][i] = tarr[j]; 
			}
		}
		return src;
	}
	

	
	protected static void butterfly(Complex a, Complex b, Complex w) {
		Complex anew = a.add(b.multy(w));
		Complex bnew = a.sub(b.multy(w));
		copyValues(a, anew);
		copyValues(b, bnew);
	}

	protected static void butterfly(Complex a00, Complex a01, Complex a10, Complex a11, Complex w1, Complex w2) {
		Complex newa00 = a00.clone();
		Complex newa10 = a00.clone();
		Complex newa01 = a00.clone();
		Complex newa11 = a00.clone();
		
		Complex t = w1.multy(a10);
		newa00 = newa00.add(t);
		newa10 = newa10.sub(t);
		newa01 = newa01.add(t);
		newa11 = newa11.sub(t);
		
		t = w2.multy(a01);
		newa00 = newa00.add(t);
		newa10 = newa10.add(t);
		newa01 = newa01.sub(t);
		newa11 = newa11.sub(t);
		
		t = w1.multy(w2.multy(a11));
		newa00 = newa00.add(t);
		newa10 = newa10.sub(t);
		newa01 = newa01.sub(t);
		newa11 = newa11.add(t);
		
		copyValues(a00, newa00);
		copyValues(a01, newa01);
		copyValues(a10, newa10);
		copyValues(a11, newa11);
	}

	protected static void copyValues(Complex a, Complex b) {
		a.setReal(b.getReal());
		a.setImg(b.getImg());
	}

	protected static void inverse(Complex[] arr){
		int dl = arr.length/2;
		int st = arr.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				Complex s0 = arr[i];
				arr[i] = arr[j];
				arr[j] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	protected static void inverse(Complex[][] src){
		int N = src.length;
		for (int i = 0; i < N; i++){
			inverse(src[i]);
		}
		int dl = src.length/2;
		int st = src.length - 1;
		int j = 0;
		for (int i = 0; i < st; i++){
			if (i < j){
				Complex[] s0 = src[i];
				src[i] = src[j];
				src[j] = s0;
			}
			int k = dl;
			while (k <= j){
				j = j - k;
				k = k >> 1;
			}
			j = j + k;
		}
	}
	
	public static void main(String[] args){
		float[][] arr = Generator.generateSimpleMatrix(8);
		Printer.printVector(arr);
		long time = System.nanoTime();
		Complex[][] res = fft2D(arr, FFTConst.SIMPLE_MODE);
		System.out.println("RESULT");
		Printer.printVector(res);
		time = System.nanoTime()-time;
		System.out.println(time);
	}
	
}
